package com.me.websocket;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

import org.java_websocket.WebSocket;

public class WsPool {
    private static final Map<WebSocket, String> wsUserMap = new HashMap<>();
    private static final ReadWriteLock rwlock = new ReentrantReadWriteLock();
    private static final Lock rlock = rwlock.readLock();
    private static final Lock wlock = rwlock.writeLock();

    // 通过websocket连接获取其对应的用户名
    public static String getUserByWs(WebSocket conn) {
        rlock.lock();
        try {
            return wsUserMap.get(conn);
        }
        finally {
            rlock.unlock();
        }
    }

    // 通过websocket连接获取其对应的IP地址
    public static String getIpByWs(WebSocket conn) {
        return conn.getRemoteSocketAddress().getAddress().toString().substring(1);
    }

    // 根据userName获取WebSocket,这是一个list,此处取第一个
    // 因为有可能多个websocket对应一个userName（但一般是只有一个，因为在close方法中，我们将失效的websocket连接去除了）
    public static WebSocket getWsByUser(String userName) {
        rlock.lock();
        try {
            Set<WebSocket> keySet = wsUserMap.keySet();
            for (WebSocket conn : keySet) {
                String user = wsUserMap.get(conn);
                if (user.equals(userName)) {
                    return conn;
                }
            }
        }
        finally {
            rlock.unlock();
        }
        return null;
    }

    // 向连接池中添加连接
    public static void addUser(String userName, WebSocket conn) {
        wlock.lock();
        try {
            wsUserMap.put(conn, userName); // 添加连接
        }
        finally {
            wlock.unlock();
        }
    }

    // 获取所有连接池中的用户，因为set是不允许重复的，所以可以得到无重复的user数组
    public static Collection<String> getOnlineUser() {
        List<String> setUsers = new ArrayList<>();
        rlock.lock();
        try {
            setUsers.addAll(wsUserMap.values());
        }
        finally {
            rlock.unlock();
        }
        return setUsers;
    }

    // 移除连接池中的连接
    public static boolean removeUser(WebSocket conn) {
        wlock.lock();
        try {
            if (wsUserMap.containsKey(conn)) {
                wsUserMap.remove(conn); // 移除连接
                return true;
            }
            else {
                return false;
            }
        }
        finally {
            wlock.unlock();
        }
    }

    // 向特定的用户发送数据
    public static void sendMessageToUser(WebSocket conn, String message) {
//        if (null != conn && null != wsUserMap.get(conn)) {
//            conn.send(message);
//        }
        if (null != conn) {
            conn.send(message);
            System.out.println("Message to " + getIpByWs(conn) + " : " + message);
        }
        else {
            System.out.println("Response nothing!");
        }
    }

    // 向所有的用户发送消息
    public static void sendMessageToAll(String message) {
        rlock.lock();
        try {
            Set<WebSocket> keySet = wsUserMap.keySet();
            for (WebSocket conn : keySet) {
                String user = wsUserMap.get(conn);
                if (user != null) {
                    conn.send(message);
                }
            }
        }
        finally {
            rlock.unlock();
        }
    }
}